# scripts/prepare_data.py
import os
import json

data_dir = "data"
jsonl_path = os.path.join(data_dir, "dataset.jsonl")

data_dir = "data/images"
baihe_dir = os.path.join(data_dir, "baihe")
dangshen_dir = os.path.join(data_dir, "dangshen")
gouqi_dir = os.path.join(data_dir, "gouqi")
huaihua_dir = os.path.join(data_dir, "huaihua")
jinyinhua_dir = os.path.join(data_dir, "jinyinhua")

# 清空dataset.jsonl文件
open(jsonl_path, 'w').close()



with open(jsonl_path, 'w', encoding='utf-8') as f:
    for img_name in os.listdir(baihe_dir):
        if img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # 示例：从文件名提取名称（你可能需要更复杂的逻辑）
            herb_name = img_name.split('.')[0].replace('herb_', '')
            text = f"这是一张中药材图片，名称是 baihe。"
            item = {"image": f"./data/images/baihe/{img_name}", "text": text}
            f.write(json.dumps(item, ensure_ascii=False) + '\n')
    for img_name in os.listdir(dangshen_dir):
        if img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # 示例：从文件名提取名称（你可能需要更复杂的逻辑）
            herb_name = img_name.split('.')[0].replace('herb_', '')
            text = f"这是一张中药材图片，名称是 dangshen。"
            item = {"image": f"./data/images/dangshen/{img_name}", "text": text}
            f.write(json.dumps(item, ensure_ascii=False) + '\n')
    for img_name in os.listdir(gouqi_dir):
        if img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # 示例：从文件名提取名称（你可能需要更复杂的逻辑）
            herb_name = img_name.split('.')[0].replace('herb_', '')
            text = f"这是一张中药材图片，名称是 gouqi。"
            item = {"image": f"./data/images/gouqi/{img_name}", "text": text}
            f.write(json.dumps(item, ensure_ascii=False) + '\n')
    for img_name in os.listdir(huaihua_dir):
        if img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # 示例：从文件名提取名称（你可能需要更复杂的逻辑）
            herb_name = img_name.split('.')[0].replace('herb_', '')
            text = f"这是一张中药材图片，名称是 huaihua。"
            item = {"image": f"./data/images/huaihua/{img_name}", "text": text}
            f.write(json.dumps(item, ensure_ascii=False) + '\n')
    for img_name in os.listdir(jinyinhua_dir):
        if img_name.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
            # 示例：从文件名提取名称（你可能需要更复杂的逻辑）
            herb_name = img_name.split('.')[0].replace('herb_', '')
            text = f"这是一张中药材图片，名称是 jinyinhua。"
            item = {"image": f"./data/images/jinyinhua/{img_name}", "text": text}
            f.write(json.dumps(item, ensure_ascii=False) + '\n')

print("✅ 预处理数据完成！")